from pprint import pprint

from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_community.tools.sql_database.tool import ListSQLDatabaseTool

from llm_utils import llm_qw3

db = SQLDatabase.from_uri('sqlite:///./chinook.db')
# print(db.get_usable_table_names())

# res = db.run('select * from artists limit 9;')
# print(res)


# 数据库自带工具
toolkits = SQLDatabaseToolkit(db=db, llm=llm_qw3)
tools = toolkits.get_tools()


"""
sql_db_query
sql_db_schema
sql_db_list_tables
sql_db_query_checker
"""
for tl in tools:
   print(f'工具名称：{tl.name}， 工具描述：{tl.description}')

# 两种方法也可以执行
tables_tool = next(t for t in tools if t.name == 'sql_db_list_tables')
print(tables_tool.invoke(''))

table_tool = ListSQLDatabaseTool(db=db)
print(table_tool.invoke(''))


# 获取表结构
tables_tool = next(t for t in tools if t.name == 'sql_db_schema')
print(tables_tool.invoke('artists'))





